"""
@author: j.l.liu
@create: 
@description: 把文本切分后变成了可检索的向量索引，用户提问，调取LLM/DeepSeek
"""
# app.py
from pathlib import Path

from fastapi import FastAPI
from langchain.chat_models import ChatOpenAI
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.llms import Ollama  # 本地开源 LLM

# 1. 加载本地知识库
db = Chroma(persist_directory="./chroma_db",
            embedding_function=HuggingFaceEmbeddings(
                model_name="./models/text2vec-base-chinese",
                model_kwargs={"local_files_only": True}))  # 不再联网))
retriever = db.as_retriever(search_kwargs={"k": 3})

#  2. 选模型：这里用 Ollama + llama3-cn，可换成 ChatGLM3-6B、OpenAI 等 --X-->模型下载安装，对本机内存要求较高
# llm = Ollama(model="llama3-cn")
# 2.1. 与 OpenAI 格式兼容 --不行-->收费
# llm = ChatOpenAI(
#     openai_api_base="https://api.deepseek.com/v1",
#     openai_api_key="sk-b20dd502f93647ec87021bc95334aeac",   # 填你的 Key
#     model_name="deepseek-chat",
#     temperature=0.7
# )
# -------------------- LLM 配置 --------------------
# 2.2. 阿里云百炼 DeepSeek --->50 万 tokens 的 DeepSeek-R1-250120 推理额度（可支撑1000 ~ 5000次简单对话。）有效期2年，用完再按量付费即可
llm = ChatOpenAI(
    openai_api_base="https://ark.cn-beijing.volces.com/api/v3",
    openai_api_key="2f357ee4-84ff-4460-9ba9-ee25d0f9dace",  # 直接填写
    # openai_api_key=os.getenv("ARK_API_KEY"),  # 建议放环境变量
    model="deepseek-r1-250120",
    temperature=0.7
)
# 3. 构建 RAG 链
qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, chain_type="stuff")

# 4. FastAPI 暴露接口
app = FastAPI(title="本地问答机器人-阿里云百炼")


@app.post("/ask")
def ask(question: str):
    return {"answer": qa_chain.invoke(question)["result"]}


if __name__ == "__main__":
    import uvicorn, dotenv

    # dotenv.load_dotenv()  # 读取 .env 里的 ARK_API_KEY

    uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True)
